import re
import warnings
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.formula.api as smf
from lifelines import KaplanMeierFitter
from statsmodels.duration.hazard_regression import PHReg
from statsmodels.tools.sm_exceptions import ConvergenceWarning

# 设置样式
plt.style.use('seaborn-v0_8-whitegrid')

# 忽略警告
warnings.simplefilter('ignore', ConvergenceWarning)
plt.switch_backend('Agg')  # 使用非交互式后端避免显示问题

# 设置中文字体显示
plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei', 'Arial Unicode MS']
plt.rcParams['axes.unicode_minus'] = False

# BMI分组标准
BMI_BINS = [20, 28, 32, 36, 40, np.inf]
BMI_LABELS = ['20-28', '28-32', '32-36', '36-40', '40+']


def parse_week(week_str):
    """将孕周从 '11w+6' 转换为小数周数"""
    if isinstance(week_str, str):
        match = re.match(r"(\d+)w\+(\d+)", week_str)
        if match:
            weeks = int(match.group(1))
            days = int(match.group(2))
            return weeks + days / 7
    return np.nan


def load_and_preprocess_data(file_path):
    """加载和预处理数据"""
    try:
        male_data = pd.read_excel(file_path, sheet_name='男胎检测数据')
    except Exception as e:
        print(f"数据加载失败: {e}")
        return None

    # 解析孕周数据
    male_data['孕周'] = male_data['检测孕周'].apply(parse_week)

    # 筛选有效数据
    male_data = male_data.dropna(subset=['Y染色体浓度', '孕妇BMI', '孕周'])
    male_data = male_data.rename(columns={
        '孕妇代码': 'subject_id',
        '孕妇BMI': 'bmi',
        '年龄': 'age',
        '身高': 'height',
        '体重': 'weight',
        'Y染色体浓度': 'Y浓度'
    })

    return male_data


def create_survival_data(df_long):
    """创建生存分析数据"""
    survival_data = []
    unique_subjects = df_long['subject_id'].unique()

    for subject in unique_subjects:
        subject_df = df_long[df_long['subject_id'] == subject].sort_values('孕周')

        # 检查是否达到阈值
        reached_threshold = subject_df[subject_df['Y浓度'] >= 0.04]

        if not reached_threshold.empty:
            # 找到首次达标时间
            event_time = reached_threshold['孕周'].min()
            status = 1  # 事件发生

            # 找到达标前最后一次检测时间
            before_event = subject_df[subject_df['孕周'] < event_time]
            prev_week = before_event['孕周'].max() if not before_event.empty else np.nan
        else:
            # 右删失情况
            event_time = subject_df['孕周'].max()
            status = 0  # 删失
            prev_week = np.nan

        # 收集协变量
        bmi_val = subject_df['bmi'].mean()
        age_val = subject_df['age'].mean()
        height_val = subject_df['height'].mean()
        weight_val = subject_df['weight'].mean()

        survival_data.append({
            'subject_id': subject,
            'event_time': event_time,
            'status': status,
            'left_interval': prev_week,
            'bmi': bmi_val,
            'age': age_val,
            'height': height_val,
            'weight': weight_val
        })

    df_surv = pd.DataFrame(survival_data)

    # 添加BMI分组
    df_surv['bmi_group'] = pd.cut(df_surv['bmi'], bins=BMI_BINS, labels=BMI_LABELS, include_lowest=True)

    # 为简化处理，将区间删失数据转换为中点值
    def approximate_T(row):
        if not pd.isna(row['left_interval']):
            return (row['left_interval'] + row['event_time']) / 2
        return row['event_time']

    df_surv['T_approx'] = df_surv.apply(approximate_T, axis=1)

    return df_surv


def plot_survival_curves(df_surv, filename):
    """绘制生存曲线"""
    plt.figure(figsize=(10, 6))

    for group in BMI_LABELS:
        group_df = df_surv[df_surv['bmi_group'] == group]
        if len(group_df) > 0:
            kmf = KaplanMeierFitter()
            kmf.fit(
                durations=group_df['T_approx'],
                event_observed=group_df['status'],
                label=f'BMI {group}'
            )
            kmf.plot_survival_function(ci_show=False)

    plt.title('各BMI组Y染色体浓度达标生存曲线', fontsize=14)
    plt.xlabel('孕周', fontsize=12)
    plt.ylabel('未达标概率', fontsize=12)
    plt.legend(title='BMI分组')
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(filename, dpi=300)
    plt.close()

    print(f"生存曲线已保存为 {filename}")


def calculate_optimal_times(df_surv, threshold=0.9):
    """计算最佳检测时点"""
    optimal_times = {}

    for group in BMI_LABELS:
        group_df = df_surv[df_surv['bmi_group'] == group]
        if len(group_df) > 0:
            kmf = KaplanMeierFitter()
            kmf.fit(durations=group_df['T_approx'], event_observed=group_df['status'])

            # 找到达到阈值的时间点
            survival_func = kmf.survival_function_
            cumulative_incidence = 1 - survival_func.iloc[:, 0]

            # 只考虑孕周大于10周的情况
            valid_times = cumulative_incidence.index[cumulative_incidence.index > 10]
            if len(valid_times) > 0:
                cumulative_incidence_valid = cumulative_incidence.loc[valid_times]
                reached_condition = cumulative_incidence_valid >= threshold

                if reached_condition.any():
                    optimal_time = cumulative_incidence_valid[reached_condition].index.min()
                    optimal_times[group] = optimal_time
                else:
                    optimal_times[group] = np.nan
            else:
                optimal_times[group] = np.nan
        else:
            optimal_times[group] = np.nan

    # 打印结果
    print("各BMI组的最佳检测时点 (达标概率 ≥ {}):".format(threshold))
    for group, time in optimal_times.items():
        if not np.isnan(time):
            print(f"BMI {group}: {time:.2f} 周")
        else:
            print(f"BMI {group}: 未达到阈值")

    return optimal_times


def fit_mixed_effects_model(df_long):
    """拟合混合效应模型"""
    print("正在拟合混合效应模型...")

    # 使用更多协变量
    model_formula = "Y浓度 ~ 孕周 + bmi + age + height + weight"

    # 尝试不同的优化方法
    try:
        model = smf.mixedlm(model_formula, df_long,
                            groups=df_long["subject_id"],
                            re_formula="~孕周")
        result = model.fit(method='bfgs', maxiter=1000)
    except:
        # 如果失败，尝试更简单的方法
        model = smf.mixedlm(model_formula, df_long,
                            groups=df_long["subject_id"])
        result = model.fit(method='nm', maxiter=1000)

    print("混合效应模型拟合完成")
    print(result.summary())

    return result


def fit_cox_model(df_surv, longitudinal_predictions):
    """拟合Cox比例风险模型"""
    print("正在拟合Cox比例风险模型...")

    # 准备协变量
    exog_vars = df_surv[['bmi', 'age', 'height', 'weight']].copy()
    exog_vars['pred_y'] = longitudinal_predictions

    # 确保没有缺失值
    complete_cases = exog_vars.notna().all(axis=1)
    exog_vars = exog_vars[complete_cases]
    event_times = df_surv['event_time'][complete_cases]
    status = df_surv['status'][complete_cases]

    # 拟合Cox模型
    ph_model = PHReg(endog=event_times,
                     exog=exog_vars,
                     status=status)
    ph_result = ph_model.fit()

    print("Cox比例风险模型拟合完成")
    print(ph_result.summary())

    return ph_result


def sensitivity_analysis(df_long, df_surv, sigma_list=[0.01, 0.05, 0.1]):
    """敏感性分析：评估测量误差对结果的影响"""
    print("正在进行敏感性分析...")

    results = []

    for sigma in sigma_list:
        print(f"分析噪声水平: σ = {sigma}")

        # 添加噪声
        noisy_long = df_long.copy()
        noise = np.random.normal(0, sigma, len(noisy_long))
        noisy_long['Y浓度'] += noise

        # 拟合混合效应模型
        try:
            noisy_model = smf.mixedlm("Y浓度 ~ 孕周 + bmi + age + height + weight",
                                      noisy_long,
                                      groups=noisy_long["subject_id"])
            noisy_result = noisy_model.fit(method='nm', maxiter=1000)

            # 预测浓度
            noisy_long['pred_y'] = noisy_result.predict(noisy_long)

            # 计算每个个体的平均预测浓度
            mean_pred_y = noisy_long.groupby('subject_id')['pred_y'].mean()

            # 更新生存数据
            noisy_surv = df_surv.copy()
            noisy_surv = noisy_surv.merge(mean_pred_y.rename('mean_pred_y'),
                                          left_on='subject_id', right_index=True, how='left')

            # 计算最佳检测时点
            optimal_times = calculate_optimal_times(noisy_surv)
            results.append((sigma, optimal_times))

        except Exception as e:
            print(f"在噪声水平 σ = {sigma} 下拟合失败: {e}")
            results.append((sigma, None))

    # 打印敏感性分析结果
    print("\n敏感性分析结果:")
    for sigma, optimal_times in results:
        if optimal_times is not None:
            print(f"噪声水平 σ = {sigma}:")
            for group, time in optimal_times.items():
                if not np.isnan(time):
                    print(f"  BMI {group}: {time:.2f} 周")
                else:
                    print(f"  BMI {group}: 未达到阈值")
        else:
            print(f"噪声水平 σ = {sigma}: 分析失败")

    return results




# 准备数据
df_long = load_and_preprocess_data("附件.xlsx")
if not df_long is None:


    print(f"数据准备完成，共有 {len(df_long)} 条记录，{df_long['subject_id'].nunique()} 个个体")

    # 创建生存数据
    df_surv = create_survival_data(df_long)
    print(f"生存数据创建完成，共有 {len(df_surv)} 个个体")

    # 绘制生存曲线
    plot_survival_curves(df_surv, 'question_three_生存曲线.png')

    # 第一阶段：拟合混合效应模型
    mixed_result = fit_mixed_effects_model(df_long)

    # 预测浓度
    df_long['pred_y'] = mixed_result.predict(df_long)

    # 计算每个个体的平均预测浓度
    mean_pred_y = df_long.groupby('subject_id')['pred_y'].mean()
    df_surv = df_surv.merge(mean_pred_y.rename('mean_pred_y'),
                            left_on='subject_id', right_index=True, how='left')

    # 第二阶段：拟合Cox比例风险模型
    cox_result = fit_cox_model(df_surv, df_surv['mean_pred_y'])

    # 计算最佳检测时点
    optimal_times = calculate_optimal_times(df_surv)

    # 敏感性分析
    sensitivity_results = sensitivity_analysis(df_long, df_surv)




